import cv2
import numpy as np
import os
import csv
import math
from scipy.fft import fft, fftfreq
from scipy import signal # 导入信号处理库
import matplotlib.pyplot as plt
import concurrent.futures
from functools import partial

# --- 全局变量用于鼠标回调 ---
line_points = []
calibration_image_display = None

# --- Matplotlib 中文和负号显示配置 ---
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


def mouse_callback(event, x, y, flags, param):
    """处理鼠标点击事件以进行标定"""
    global line_points, calibration_image_display
    
    # 确保 calibration_image_display 不是 None
    if calibration_image_display is None:
        return

    if event == cv2.EVENT_LBUTTONDOWN:
        if len(line_points) < 2:
            line_points.append((x, y))
            cv2.circle(calibration_image_display, (x, y), 5, (0, 255, 0), -1)
            if len(line_points) == 2:
                cv2.line(calibration_image_display, line_points[0], line_points[1], (0, 255, 0), 2)
            cv2.imshow("标定窗口 - 请画一条线", calibration_image_display)

def get_pixel_to_mm_ratio(image_path):
    """通过用户交互获取像素到毫米的转换比例。"""
    global line_points, calibration_image_display
    line_points = []

    image = cv2.imread(image_path)
    if image is None:
        print(f"错误：无法读取用于标定的图像 {image_path}")
        return None

    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    
    # --- 优化：增加高斯模糊来稳定轮廓 ---
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    
    # 计算初始阈值
    initial_thresh, _ = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    
    # 创建一个窗口用于调整阈值
    cv2.namedWindow("阈值调整窗口")
    thresh_value = initial_thresh
    cv2.createTrackbar("阈值", "阈值调整窗口", int(thresh_value), 255, lambda x: None)
    
    print("\n--- 阈值调整步骤 ---")
    print("1. 在弹出的 '阈值调整窗口' 中，使用滑块调整二值化阈值，直到满意为止。")
    print("2. 按 'q' 键确认并继续。")
    
    binary = None
    while True:
        thresh_value = cv2.getTrackbarPos("阈值", "阈值调整窗口")
        _, binary = cv2.threshold(blurred, thresh_value, 255, cv2.THRESH_BINARY_INV)
        display_image = cv2.cvtColor(binary, cv2.COLOR_GRAY2BGR)
        cv2.imshow("阈值调整窗口", display_image)
        if cv2.waitKey(20) & 0xFF == ord('q'):
            break
    
    cv2.destroyWindow("阈值调整窗口")
    
    # 进行形态学操作以优化二值化图像
    kernel = np.ones((5, 5), np.uint8)
    binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel, iterations=2)
    binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel, iterations=3)
    
    # 查找轮廓
    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    calibration_image_display = cv2.cvtColor(binary, cv2.COLOR_GRAY2BGR)
    
    # 绘制最大轮廓
    if contours:
        largest_contour = max(contours, key=cv2.contourArea)
        cv2.drawContours(calibration_image_display, [largest_contour], -1, (0, 255, 0), 2)

    cv2.namedWindow("标定窗口 - 请画一条线")
    cv2.setMouseCallback("标定窗口 - 请画一条线", mouse_callback)

    print("\n--- 全局标定步骤 (仅需一次) ---")
    print("1. 在弹出的 '标定窗口' 中，用鼠标左键点击两次，画一条你知道其实际长度的直线。")
    print("2. 此标定将应用于所有找到的实验文件夹。")

    while len(line_points) < 2:
        cv2.imshow("标定窗口 - 请画一条线", calibration_image_display)
        if cv2.waitKey(20) & 0xFF == ord('q'):
            cv2.destroyAllWindows()
            print("标定被用户取消。")
            return None

    try:
        real_length_mm = float(input("请输入您所画直线的实际长度 (单位: 毫米)，然后按 Enter: "))
    except ValueError:
        print("输入无效，请输入一个数字。")
        cv2.destroyAllWindows()
        return None

    cv2.destroyAllWindows()

    p1, p2 = line_points
    pixel_distance = math.sqrt((p2[0] - p1[0])**2 + (p2[1] - p1[1])**2)

    if real_length_mm <= 0:
        print("实际长度必须大于0。")
        return None

    pixels_per_mm = pixel_distance / real_length_mm
    print(f"标定完成: 像素距离 = {pixel_distance:.2f} px, 实际长度 = {real_length_mm} mm")
    print(f"计算出的全局比例尺: {pixels_per_mm:.2f} 像素/毫米")
    
    return pixels_per_mm

def analyze_and_plot_vibration(csv_path, pixels_per_mm, frame_rate, output_folder):
    """读取质心数据，进行全面的振动分析并绘图，同时保存数据到CSV文件。"""
    x_coords_pixel, y_coords_pixel = [], []
    
    try:
        with open(csv_path, 'r', encoding='utf-8') as csvfile:
            reader = csv.reader(csvfile)
            next(reader) # 跳过表头
            for row in reader:
                x_coords_pixel.append(float(row[1]))
                y_coords_pixel.append(float(row[2]))
    except (StopIteration, FileNotFoundError):
        print(f"  - 警告: 无法读取或找不到CSV文件 {csv_path}。跳过分析。")
        return
    except Exception as e:
        print(f"  - 错误: 读取CSV文件时出错: {e}。跳过分析。")
        return

    if not y_coords_pixel:
        print(f"  - 警告: {csv_path} 中无数据。跳过分析。")
        return

    x_coords_pixel = np.array(x_coords_pixel)
    y_coords_pixel = np.array(y_coords_pixel)

    # 计算相对于平均位置的位移
    disp_x_mm = (x_coords_pixel - np.mean(x_coords_pixel)) / pixels_per_mm
    disp_y_mm = (y_coords_pixel - np.mean(y_coords_pixel)) / pixels_per_mm
    
    N = len(disp_y_mm)
    T = 1.0 / frame_rate
    time = np.linspace(0.0, N * T, N, endpoint=False)
    
    # 保存位移数据到CSV文件
    displacement_csv_path = os.path.join(output_folder, 'displacement_data.csv')
    with open(displacement_csv_path, 'w', newline='', encoding='utf-8') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(['时间(s)', 'X方向位移(mm)', 'Y方向位移(mm)'])
        for t, x, y in zip(time, disp_x_mm, disp_y_mm):
            writer.writerow([t, x, y])
    print(f"  - 位移数据已保存到: {displacement_csv_path}")
    
    # 计算并保存运动指标
    motion_metrics_csv = os.path.join(output_folder, 'motion_metrics.csv')
    max_distance = 0.0
    max_deviation = 0.0
    
    # 计算初始位置（第一个点）
    initial_x, initial_y = disp_x_mm[0], disp_y_mm[0]
    
    # 计算最大移动距离（轨迹总长度）
    total_distance = 0.0
    prev_x, prev_y = disp_x_mm[0], disp_y_mm[0]
    for i in range(1, len(disp_x_mm)):
        dx = disp_x_mm[i] - prev_x
        dy = disp_y_mm[i] - prev_y
        segment_length = math.sqrt(dx**2 + dy**2)
        total_distance += segment_length
        prev_x, prev_y = disp_x_mm[i], disp_y_mm[i]
        if segment_length > max_distance:
            max_distance = segment_length
    
    # 计算最大偏离距离（各点到初始位置的距离）
    for i in range(len(disp_x_mm)):
        dx = disp_x_mm[i] - initial_x
        dy = disp_y_mm[i] - initial_y
        distance = math.sqrt(dx**2 + dy**2)
        if distance > max_deviation:
            max_deviation = distance
    
    # 保存运动指标
    with open(motion_metrics_csv, 'w', newline='', encoding='utf-8') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(['指标名称', '值(mm)', '描述'])
        writer.writerow(['最大单步移动距离', f'{max_distance:.4f}', '连续两帧间最大移动距离'])
        writer.writerow(['最大偏离中心距离', f'{max_deviation:.4f}', '偏离初始位置的最大距离'])
        writer.writerow(['运动轨迹总长度', f'{total_distance:.4f}', '质心运动轨迹总长度'])
    
    print(f"  - 运动指标数据已保存到: {motion_metrics_csv}")
    print(f"  - 最大单步移动距离: {max_distance:.4f} mm")
    print(f"  - 最大偏离中心距离: {max_deviation:.4f} mm")
    print(f"  - 运动轨迹总长度: {total_distance:.4f} mm")

    # 1. 质心运动轨迹图
    plt.figure(figsize=(8, 8))
    plt.plot(disp_x_mm, disp_y_mm, linewidth=0.5)
    plt.title('质心运动轨迹')
    plt.xlabel('X方向位移 (mm)')
    plt.ylabel('Y方向位移 (mm)')
    plt.grid(True)
    plt.axis('equal')
    plt.savefig(os.path.join(output_folder, 'centroid_trajectory.png'))
    plt.close()

    # 2. X, Y方向位移-时间图
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 10), sharex=True)
    ax1.plot(time, disp_y_mm)
    ax1.set_title('Y方向振动位移-时间图')
    ax1.set_ylabel('位移 (mm)')
    ax1.grid(True)
    ax2.plot(time, disp_x_mm)
    ax2.set_title('X方向振动位移-时间图')
    ax2.set_ylabel('位移 (mm)')
    ax2.set_xlabel('时间 (s)')
    ax2.grid(True)
    plt.savefig(os.path.join(output_folder, 'displacement_vs_time_XY.png'))
    plt.close()

    def plot_fft_and_find_peaks(disp_data, axis_name):
        yf = fft(disp_data)
        xf = fftfreq(N, T)[:N//2]
        amplitude = 2.0/N * np.abs(yf[0:N//2])

        plt.figure(figsize=(12, 6))
        plt.plot(xf, amplitude)
        plt.title(f'{axis_name}方向频域分析 (FFT)')
        plt.xlabel('频率 (Hz)')
        plt.ylabel('振幅 (mm)')
        plt.grid(True)
        
        peaks, _ = signal.find_peaks(amplitude, prominence=np.max(amplitude)/10, height=np.max(amplitude)/5)
        top_peaks_indices = sorted(peaks, key=lambda p: amplitude[p], reverse=True)[:3]
        
        with open(os.path.join(output_folder, 'frequency_results.txt'), 'a', encoding='utf-8') as f:
            f.write(f"\n--- {axis_name}方向前三阶固有频率 ---\n")
            if not top_peaks_indices:
                f.write("未能找到显著的频率峰值。\n")
            else:
                for i, peak_idx in enumerate(top_peaks_indices):
                    peak_freq = xf[peak_idx]
                    peak_amp = amplitude[peak_idx]
                    result_line = f"  第 {i+1} 阶频率: {peak_freq:.2f} Hz (振幅: {peak_amp:.4f} mm)\n"
                    f.write(result_line)
                    plt.axvline(peak_freq, color='r', linestyle='--', alpha=0.7)
                    plt.text(peak_freq, amplitude[peak_idx], f' {peak_freq:.2f} Hz', color='red')
        
        # 保存FFT数据到CSV文件
        fft_csv_path = os.path.join(output_folder, f'fft_data_{axis_name}.csv')
        with open(fft_csv_path, 'w', newline='', encoding='utf-8') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow(['频率(Hz)', '振幅(mm)'])
            for freq, amp in zip(xf, amplitude):
                writer.writerow([freq, amp])
        print(f"  - {axis_name}方向FFT数据已保存到: {fft_csv_path}")
        
        plt.savefig(os.path.join(output_folder, f'frequency_analysis_{axis_name}.png'))
        plt.close()

    def plot_psd(disp_data, axis_name):
        freqs, psd = signal.welch(disp_data, fs=frame_rate, nperseg=1024)
        plt.figure(figsize=(12, 6))
        plt.semilogy(freqs, psd)
        plt.title(f'{axis_name}方向功率谱密度 (PSD)')
        plt.xlabel('频率 (Hz)')
        plt.ylabel('功率/频率 (dB/Hz)')
        plt.grid(True)
        
        # 保存PSD数据到CSV文件
        psd_csv_path = os.path.join(output_folder, f'psd_data_{axis_name}.csv')
        with open(psd_csv_path, 'w', newline='', encoding='utf-8') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow(['频率(Hz)', '功率/频率(dB/Hz)'])
            for freq, p in zip(freqs, psd):
                writer.writerow([freq, p])
        print(f"  - {axis_name}方向PSD数据已保存到: {psd_csv_path}")
        
        plt.savefig(os.path.join(output_folder, f'psd_analysis_{axis_name}.png'))
        plt.close()

    # 清空之前的频率结果文件
    freq_results_path = os.path.join(output_folder, 'frequency_results.txt')
    if os.path.exists(freq_results_path):
        os.remove(freq_results_path)

    plot_fft_and_find_peaks(disp_y_mm, 'Y')
    plot_fft_and_find_peaks(disp_x_mm, 'X')
    plot_psd(disp_y_mm, 'Y')
    plot_psd(disp_x_mm, 'X')


def process_single_image(file_info, input_folder, output_folder):
    """处理单个图像文件的函数，设计为在线程池中运行。"""
    i, filename, should_save = file_info
    file_path = os.path.join(input_folder, filename)
    image = cv2.imread(file_path, cv2.IMREAD_UNCHANGED)

    if image is None: return None

    # 确保总是处理灰度图
    if len(image.shape) == 2:
        gray_image = image
    else: # len(image.shape) == 3
        gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
    # --- 关键优化 1: 在二值化前进行高斯模糊以减少噪声，稳定边缘 ---
    blurred_image = cv2.GaussianBlur(gray_image, (5, 5), 0)
        
    # 使用Otsu方法进行二值化
    _, binary_image = cv2.threshold(blurred_image, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    
    # 形态学操作
    kernel = np.ones((5, 5), np.uint8)
    binary_image = cv2.morphologyEx(binary_image, cv2.MORPH_OPEN, kernel, iterations=2)
    binary_image = cv2.morphologyEx(binary_image, cv2.MORPH_CLOSE, kernel, iterations=3)
    
    # 查找轮廓
    contours, _ = cv2.findContours(binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    if contours:
        largest_contour = max(contours, key=cv2.contourArea)
        M = cv2.moments(largest_contour)
        
        if M["m00"] != 0:
            # --- 关键优化 2: 保持质心坐标为浮点数，以实现亚像素精度 ---
            cX_float = M["m10"] / M["m00"]
            cY_float = M["m01"] / M["m00"]
        else:
            # 如果找不到矩，使用边界框中心作为备用方案
            x, y, w, h = cv2.boundingRect(largest_contour)
            cX_float, cY_float = float(x + w // 2), float(y + h // 2)

        # 如果需要保存带标记的图像
        if should_save:
            # 为了绘制，需要将原始图像转为彩色
            if len(image.shape) == 2:
                image_for_drawing = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
            else:
                image_for_drawing = image.copy()
            
            # 为了在图像上绘制，我们需要整数坐标
            draw_cX, draw_cY = int(round(cX_float)), int(round(cY_float))
            
            # 绘制轮廓和质心
            cv2.drawContours(image_for_drawing, [largest_contour], -1, (0, 255, 0), 2)
            cv2.circle(image_for_drawing, (draw_cX, draw_cY), 7, (0, 0, 255), -1)
            cv2.putText(image_for_drawing, f"({draw_cX}, {draw_cY})", (draw_cX + 10, draw_cY + 10),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)
            output_path = os.path.join(output_folder, filename)
            cv2.imwrite(output_path, image_for_drawing)
        
        # 返回浮点数坐标用于精确分析
        return (filename, cX_float, cY_float)
        
    return None

def process_images_in_folder(input_folder, output_folder, num_images_to_save):
    """使用线程池加速处理文件夹中的所有图像。"""
    try:
        file_list = sorted([f for f in os.listdir(input_folder) if f.lower().endswith((".bmp", ".png", ".jpg", ".jpeg"))])
        if not file_list:
            print(f"  - 警告: 文件夹 {input_folder} 中没有找到图像文件。")
            return None
    except FileNotFoundError:
        print(f"  - 错误: 找不到文件夹 {input_folder}。")
        return None

    total_files = len(file_list)
    save_indices = set(np.linspace(0, total_files - 1, num_images_to_save, dtype=int))
    tasks = [(i, filename, i in save_indices) for i, filename in enumerate(file_list)]
    
    worker_func = partial(process_single_image, input_folder=input_folder, output_folder=output_folder)
    
    results = []
    # 使用tqdm库显示进度条（如果已安装）
    try:
        from tqdm import tqdm
        with concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count()) as executor:
            results = list(tqdm(executor.map(worker_func, tasks), total=total_files, desc=f"处理 {os.path.basename(input_folder)}"))
    except ImportError:
        print(f"正在处理文件夹: {input_folder} (共 {total_files} 张图像)...")
        with concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count()) as executor:
            results = list(executor.map(worker_func, tasks))
        print("处理完成。")

    # 过滤掉处理失败的结果
    centroids_data = [res for res in results if res is not None]
    
    if not centroids_data:
        print(f"  - 警告: 未能从 {input_folder} 的任何图像中提取质心。")
        return None

    csv_path = os.path.join(output_folder, 'centroids.csv')
    try:
        with open(csv_path, 'w', newline='', encoding='utf-8') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow(['文件名', '质心_X', '质心_Y'])
            writer.writerows(centroids_data)
        return csv_path
    except IOError as e:
        print(f"  - 错误: 无法写入CSV文件 {csv_path}: {e}")
        return None

# --- 主程序 ---
if __name__ == "__main__":
    # 1. 获取用户输入的根目录
    root_input_dir = input("请输入包含所有实验工况的根文件夹地址: ")
    if not os.path.isdir(root_input_dir):
        print(f"错误: 文件夹 '{root_input_dir}' 不存在。")
        exit()

    # 2. 查找所有包含图像文件的 "叶子" 文件夹
    experiment_folders = []
    image_extensions = {'.bmp', '.png', '.jpg', '.jpeg'}
    for dirpath, _, filenames in os.walk(root_input_dir):
        if any(os.path.splitext(f)[1].lower() in image_extensions for f in filenames):
            # 检查这是否是一个 "叶子" 节点（不包含其他也含图像的子目录）
            subdirs_with_images = False
            for sub_dir in os.listdir(dirpath):
                sub_dir_path = os.path.join(dirpath, sub_dir)
                if os.path.isdir(sub_dir_path):
                    for _, _, sub_files in os.walk(sub_dir_path):
                        if any(os.path.splitext(f)[1].lower() in image_extensions for f in sub_files):
                            subdirs_with_images = True
                            break
                    if subdirs_with_images:
                        break
            if not subdirs_with_images:
                 experiment_folders.append(dirpath)


    if not experiment_folders:
        print(f"在 '{root_input_dir}' 及其子文件夹中未找到任何包含图像文件的叶子文件夹。")
        exit()
    
    print(f"\n找到了 {len(experiment_folders)} 个实验文件夹，准备处理...")
    for folder in experiment_folders:
        print(f"  - {folder}")


    # 3. 进行一次性全局标定
    first_folder = experiment_folders[0]
    first_image_name = next((f for f in sorted(os.listdir(first_folder)) if os.path.splitext(f)[1].lower() in image_extensions), None)
    
    if not first_image_name:
        print("错误：在第一个实验文件夹中找不到任何图像用于标定。")
        exit()

    first_image_path = os.path.join(first_folder, first_image_name)
    pixels_per_mm = get_pixel_to_mm_ratio(first_image_path)

    if pixels_per_mm is None:
        print("标定失败，程序退出。")
        exit()

    # 4. 创建总输出目录
    root_output_dir = os.path.join(os.path.dirname(root_input_dir.rstrip('/\\')), os.path.basename(root_input_dir.rstrip('/\\')) + "_processed")
    os.makedirs(root_output_dir, exist_ok=True)
    print(f"\n所有结果将保存在: {root_output_dir}")

    # 5. 循环处理每个实验文件夹
    FRAME_RATE = 1000  # 假设帧率为1000Hz，可根据实际情况调整
    NUM_IMAGES_TO_SAVE = 10 # 每个工况保存10张带标记的图片

    for i, input_folder in enumerate(experiment_folders):
        relative_path = os.path.relpath(input_folder, root_input_dir)
        output_folder = os.path.join(root_output_dir, relative_path)
        os.makedirs(output_folder, exist_ok=True)

        print(f"\n--- [{i+1}/{len(experiment_folders)}] 开始处理实验: {relative_path} ---")
        
        # 处理图像并获取质心数据
        centroids_csv_file = process_images_in_folder(input_folder, output_folder, NUM_IMAGES_TO_SAVE)
        
        # 进行振动分析和绘图
        if centroids_csv_file:
            print(f"  - 质心数据已保存，开始振动分析...")
            analyze_and_plot_vibration(centroids_csv_file, pixels_per_mm, FRAME_RATE, output_folder)
            print(f"  - 分析完成，结果保存在: {output_folder}")
        else:
            print(f"  - 警告: 未能生成质心数据CSV文件，跳过此文件夹的分析。")

    print("\n\n所有实验工况均已处理完毕！")
